import os
from unstructured.partition.pdf import partition_pdf
from unstructured.staging.base import elements_to_json
from langchain.chat_models import ChatOpenAI
from langchain.schema.messages import HumanMessage
from dotenv import load_dotenv
from PIL import Image
import base64
import io
from typing import Any
from pydantic import BaseModel

###pdf处理，表格
filename = "/1.pdf"
elements = partition_pdf(
    filename=filename,
    # 使用高分辨率提取策略，可以准确的识别文档结构
    strategy="hi_res",
    # 自动识别pdf中的表格结构
    infer_table_structure=True,
    # 使用yolox模型进行元素的识别
    model_name="yolox"
)

elements[-4].metadata.text_as_html  # 将表格内容转化成html

# pdf处理，图片
# 使用多模态模型对图像生成摘要，对摘要进行嵌入
# 获取图像的嵌入，比如使用CLIP模型对文本和图像同时做嵌入
# 直接用图像本身
# pip install unstructured[all-docs]


raw_pdf_elements = partition_pdf(
    filename="C:\\download\\1.pdf",
    # 提取pdf中的嵌入图像功能，解析的图像会被保存至指定的目录中
    extract_images_in_pdf=True,
    # 使用yolox来识别pdf中的表格，并提取其结构和内容。同时，模型还会尝试识别文档中的标题，这些标题可以作为后续分块的依据
    infer_table_structure=True,
    # 按标题分块。pdf中的每个子章节（标题）会被视为一个独立的块
    chunking_strategy="by_title",
    # chunking params to aggregate text blocks attempt to create a new chunk 3800 chars
    max_characters=4000,
    # 当某个块的字符数接近3800是，尝试创建一个新的块，这有助于避免单个块过大
    new_after_n_chars=3800,
    # 如果某些小块的字符数少于2000，则尝试将它们合并到相邻的块中，这可以减少过小的块数量，保持块的大小合理
    combine_text_under_n_chars=2000,
    image_output_dir_path="/static/pdfImages/",
)

# 生成图片摘要


load_dotenv()
llm = ChatOpenAI(model_name="gpt-4-vision-preview", temperature=0)


# function to convert image file to base64 string
def image_to_base64(image_path: str):
    with Image.open(image_path) as image_file:
        buffered = io.BytesIO()
        image_file.save(buffered, format=image_file.format)
        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
        return img_str
    image_str = image_to_base64("/static/pdfimages/figure-15-6.jpg")
    chat = ChatOpenAI(model_name="gpt-4-vision-preview", max_tokens=1024)
    msg = chat.invoke(
        [
            HumanMessage(
                content=[
                    {"type": "text", "text": "Please give a summary of the image provided."},
                    {
                        "type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_str}"},
                    },
                ]
            )
        ]
    )
    msg.content
